package  gateway.filter;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.collection.SpliteratorUtil;
import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson.JSONObject;
import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.NoArgsConstructor;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.reactivestreams.Publisher;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.NettyWriteResponseFilter;
import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.core.io.buffer.DefaultDataBufferFactory;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.io.StringReader;
import java.nio.charset.StandardCharsets;
import java.util.List;

/**
 * @Author: meng
 * @Description: 自定义返回体 - 借鉴原生类ModifyRequestBodyGatewayFilterFactory实现
 * https://docs.spring.io/spring-cloud-gateway/docs/current/reference/html/#the
 * -modifyresponsebody-gatewayfilter-factory
 * @Date: 2023/4/18 13:37
 * @Version: 1.0
 */
@Slf4j
@Component
public class LogGatewayFilterFactory extends AbstractGatewayFilterFactory<LogGatewayFilterFactory.Config> {

    public LogGatewayFilterFactory() {
        super(Config.class);
    }

    @Data
    @AllArgsConstructor
    @NoArgsConstructor
    public static class Config {

        // 不需要自定义的接口
        List<String> pathExclude;

    }

    @Override
    public GatewayFilter apply(Config config) {
        RewriteResponseGatewayFilter rewriteResponseGatewayFilter = new RewriteResponseGatewayFilter(config);
        return rewriteResponseGatewayFilter;
    }

    public class RewriteResponseGatewayFilter implements GatewayFilter, Ordered {

        private Config config;

        public RewriteResponseGatewayFilter(Config config) {
            this.config = config;
        }

        @Override
        public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {

            ServerHttpResponse originalResponse = exchange.getResponse();
            DataBufferFactory bufferFactory = originalResponse.bufferFactory();
            try {
                ServerHttpResponseDecorator decoratedResponse = new ServerHttpResponseDecorator(originalResponse) {
                    @Override
                    public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
                        if (body instanceof Flux) {
                            Flux<? extends DataBuffer> fluxBody = Flux.from(body);
                            return super.writeWith(fluxBody.buffer().map(dataBuffers -> {
                                byte[] newContent = new byte[0];
                                try {
                                    DataBufferFactory dataBufferFactory = new DefaultDataBufferFactory();
                                    DataBuffer join = dataBufferFactory.join(dataBuffers);
                                    byte[] content = new byte[join.readableByteCount()];
                                    join.read(content);
                                    DataBufferUtils.release(join);
                                    // 获取响应数据
                                    String responseStr = new String(content, StandardCharsets.UTF_8);
                                    log.info("原始响应数据responseStr:{}",newContent);
                                    // 修改响应数据
                                    String[] split = StringUtils.tokenizeToStringArray(responseStr, "\n");
                                    StringBuilder sb = new StringBuilder();
                                    for (String s : split) {
                                        String token = s;
                                        if (!token.startsWith("#")) {
                                            token = "test_" + token;
                                        }
                                        sb.append(token).append("\n");
                                    }
                                    newContent = sb.toString().getBytes(StandardCharsets.UTF_8);
                                    log.info("修改响应数据 responseStr:{}",responseStr);
                                    originalResponse.getHeaders().setContentLength(newContent.length);
                                }
                                catch (Exception e) {
                                    throw new RuntimeException(e);
                                }
                                return bufferFactory.wrap(newContent);
                            }));
                        }
                        return super.writeWith(body);
                    }

                    @Override
                    public Mono<Void> writeAndFlushWith(Publisher<? extends Publisher<? extends DataBuffer>> body) {
                        return writeWith(Flux.from(body).flatMapSequential(p -> p));
                    }
                };
                return chain.filter(exchange.mutate().response(decoratedResponse).build());
            }
            catch (Exception e) {
                log.error("RewriteResponse error:{}",e.getMessage(), e);
                return Mono.error(new Exception("RewriteResponse fail", e));
            }
        }

        @Override
        public int getOrder() {
            return NettyWriteResponseFilter.WRITE_RESPONSE_FILTER_ORDER - 2;
        }

    }

    @SneakyThrows
    public static Mono<Void> buildResponse(ServerWebExchange exchange, int code, String message) {
        ServerHttpResponse response = exchange.getResponse();
        response.setStatusCode(HttpStatus.OK);
        response.getHeaders().add(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON_UTF8_VALUE);
        JSONObject jsonObject = new JSONObject();
        jsonObject.put("code", code);
        jsonObject.put("message", message);
        jsonObject.put("data", "");
        byte[] bytes = jsonObject.toJSONString().getBytes("UTF-8");
        DataBuffer bodyDataBuffer = response.bufferFactory().wrap(bytes);
        return response.writeWith(Mono.just(bodyDataBuffer));
    }

}
